%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%              Conditional Distribution of Product Age                    %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [Phi, avec, Pmatrix, p0] = Conditional_Dist(p)

% ---------------------------- Description --------------------------------

    % This function computes the conditional distribution of product age
    % given firm age and the number of products that the firm has. This
    % distribution is referred to as Phi_{A,n}(a).

    % To obtain this function, we solve for the numeric probabilities
    % p_A(n), which allows us to get Lambda_A(n). Then, we solve the
    % differential equation for X_{A,n}(a) and manipulate to get Phi.

    % For an overview of this function, please see Section (RP-2.4.4).

% -------------------------------------------------------------------------
%                    1. Closed form Probabilities
% -------------------------------------------------------------------------

    % Determines the number of points in the age grid
    TimePeriods = p.maxTime/p.periodlength + 1;

    % Creates a linear grid for age - serves for product and firm age
    avec = linspace(0, p.maxTime, TimePeriods);

    % Computes the implied step in age grid
    delta_a = p.periodlength;

    % Computes the Gamma function over the age grid - see Equation (RP-3)
    GammaFunc = p.xstar*(1-exp(p.psi*avec))./(p.xstar*(1-exp(p.psi*avec)) - p.psi);

    % Compute p_0(a) and p_1(a) over the age grid - see Equations (RP-2)
    p0 = (p.tau+p.d0)/p.xstar*GammaFunc;
    p1 = (1-p0).*(1-GammaFunc);

    % Renames the maximum number of products for convenience
    nmax = p.maxprods;

    % Matrix for the conditional product distribution
    Pmatrix = [p1; zeros(nmax-1, TimePeriods)];

    % For each potential number of products...
    for i = 2:p.maxprods

        % ...recursively compute p_n(a) based on p_{n-1}(a) - see Equations (RP-2)
        Pmatrix(i,:) = Pmatrix(i-1,:) .* GammaFunc;

    end
        
% -------------------------------------------------------------------------
%                    2. Numeric Probabilities
% -------------------------------------------------------------------------

    % Normalize the number of age 0 firms to one
    Fo = 1;

    % Preallocate matrix for numeric probabilites - need to consider 0 products here
    LambdaMat_num = zeros(nmax+1,TimePeriods);
    
    % Boundary Condition: p_{A=0}(n=1) = 1 and p_{A=0}(n!=1) = 0
    LambdaMat_num(2,1) = Fo;

    % For each value of firm age...
    for A=2:TimePeriods

        % First, compute derivatives for n = 0 - see Equation (RP-28)
        dlda = (p.tau+p.d0)*LambdaMat_num(2,A-1);
        LambdaMat_num(1,A) = LambdaMat_num(1,A-1) + dlda*delta_a;

        % Second, compute derivatives for n = 1 - see Equation (RP-28)
        dlda = -(p.tau+p.d0+p.xstar)*LambdaMat_num(2,A-1)+(p.tau+p.d0)*2*LambdaMat_num(3,A-1);
        LambdaMat_num(2,A) = LambdaMat_num(2,A-1) + dlda*delta_a;

        % Third, compute derivatives for 1<n<nmax - see Equation (RP-28)
        Naux_grid = (3:p.maxprods)';
        dlda = -(p.tau+p.d0+p.xstar)*(Naux_grid-1).*LambdaMat_num(Naux_grid,A-1)+(p.tau+p.d0)*Naux_grid.*LambdaMat_num(Naux_grid+1,A-1);
        dlda = dlda + p.xstar*(Naux_grid-2).*LambdaMat_num(Naux_grid-1,A-1);
        LambdaMat_num(Naux_grid,A) = LambdaMat_num(Naux_grid,A-1) + dlda*delta_a;

        % Finally, compute derivatives for nmax - see Equation (RP-28)
        dlda = -(p.tau+p.d0+p.xstar)*(nmax)*LambdaMat_num(nmax+1,A-1)+(p.tau+p.d0)*(nmax+1)*LambdaMat_num(nmax+1,A-1);
        dlda = dlda + p.xstar*(nmax-1)*LambdaMat_num(nmax,A-1);  
        LambdaMat_num(nmax+1,A)= LambdaMat_num(nmax+1,A-1) + dlda*delta_a;
    
    end

    % Creates a vector with the number of products for multiplication
    prodvec = 1:nmax;   

    % Compute the matrix Lambda times the number of products
    LambdaMat_num = LambdaMat_num(2:nmax+1,:);
    LambdaMat_num_N = repmat(prodvec',1,TimePeriods).*LambdaMat_num;
        
    % Change the dimensions of LambdaN to match the dimensions of X
    LambdaN_cube_num = repmat(LambdaMat_num_N,[1 1 TimePeriods]);
    LambdaN_cube_num = permute(LambdaN_cube_num,[3 2 1]);

%--------------------------------------------------------------------------
%                 3. Differential Equations for X_{A,n}(a)
%--------------------------------------------------------------------------

    % Matrix with dimensions age (prods.), age (firm) and number of products
    Xmatrix = ones(TimePeriods, TimePeriods, p.maxprods);

    % Sets our boundary condition for the function X_{A,n}(a) - see Section (RP-2.4.4)
    for n = 1:p.maxprods
        Xmatrix(:,:,n) = tril(Xmatrix(:,:,n)).*LambdaN_cube_num(:,:,n);
    end

    % Construct objects related to the product grid that assist computation
    Naux_grid = 2:nmax-1;
    Naux_mat = permute(Naux_grid',[3 2 1]);

    % For each value of firm age...
    for A = 3:TimePeriods

            % Solve simplified Differential Equation for first row - see Section (RP-2.4.4)
            Naux = 2:nmax;
            dx = p.xstar*(Naux'- 1).*LambdaMat_num(Naux-1,A-2); 
            Xmatrix(1,A-1,Naux) = 0 + dx*delta_a;  
 
            % Construct objects related to the age grid that assist computation
            Aaux = 2:(A-1);

            % Solves complete differential equation for matrix associated with n = 1 - see Equation (RP-29)
            dx = -(p.tau+p.d0+p.xstar)*1*Xmatrix(Aaux-1, A-1, 1) + (p.tau+p.d0)*1*Xmatrix(Aaux-1, A-1, 2);
            Xmatrix(Aaux, A, 1) = Xmatrix(Aaux-1, A-1, 1) + dx*delta_a;

            % Solves complete differential equation for matrix associated with n = 2 to n = nmax-1 - see Equation (RP-29)
            Lambda_aux = permute(LambdaMat_num(Naux_grid-1,A-1),[3 2 1]);
            dx = -(p.tau+p.d0+p.xstar)*Naux_mat.*Xmatrix(Aaux-1, A-1, Naux_grid)+(p.tau+p.d0)*(Naux_mat).*Xmatrix(Aaux-1,A-1,Naux_grid+1);
            dx = dx + p.xstar*(Naux_mat-1).*(Lambda_aux+Xmatrix(Aaux-1,A-1,Naux_grid-1));
            Xmatrix(Aaux, A, Naux_grid) = Xmatrix(Aaux-1, A-1, Naux_grid) + dx*delta_a;

            % Solves complete differential equation for matrix associated with n = nmax - see Equation (RP-29)
            dx = -(p.tau+p.d0+p.xstar)*nmax*Xmatrix(Aaux-1, A-1, nmax) + (p.tau+p.d0)*nmax*Xmatrix(Aaux-1, A-1, nmax);
            dx = dx + p.xstar*(nmax-1)*(LambdaMat_num(nmax-1,A-1)+Xmatrix(Aaux-1,A-1,nmax-1));
            Xmatrix(Aaux, A, nmax) = Xmatrix(Aaux-1, A-1, nmax) + dx*delta_a;

    end

%--------------------------------------------------------------------------
%                 4. Recovering the distribution Phi
%--------------------------------------------------------------------------
    
    % To avoid dividing by zero, we divide by something small: look for Infs
    LambdaN_cube_num(LambdaN_cube_num==0) = 10e-60;

    % Obtain the Conditional Distribution by dividing by Lambda
    Phi = Xmatrix./LambdaN_cube_num;

end

